package com.lagou.mr.output;

import org.apache.hadoop.fs.FSDataOutputStream;
import org.apache.hadoop.io.IOUtils;
import org.apache.hadoop.io.NullWritable;
import org.apache.hadoop.io.Text;
import org.apache.hadoop.mapreduce.RecordWriter;
import org.apache.hadoop.mapreduce.TaskAttemptContext;

import java.io.IOException;

public class CustomWriter extends RecordWriter<Text, NullWritable> {
    private FSDataOutputStream lgOut;
    private FSDataOutputStream otherOut;

    public CustomWriter(FSDataOutputStream lgOut, FSDataOutputStream otherOut) {
        this.lgOut = lgOut;
        this.otherOut = otherOut;
    }

    @Override
    public void write(Text k, NullWritable v) throws IOException, InterruptedException {
        // 判断是否包含“lagou”输出到不同文件
        if( k.toString().contains("lagou") ){
            lgOut.write(k.toString().getBytes());
            lgOut.write("\r\n".getBytes());
        }else{
            otherOut.write(k.toString().getBytes());
            otherOut.write("\r\n".getBytes());
        }
    }

    @Override
    public void close(TaskAttemptContext context) throws IOException, InterruptedException {
        IOUtils.closeStream(lgOut);
        IOUtils.closeStream(otherOut);
    }
}
